#include "cuda_img_resize.cuh"
//#define _DEBUG

#define CUDA_IMG_RESIZE_THREAD_NUM	1024
#define CUDA_IMG_RESIZE_WARP_SIZE	32

// image.h
typedef struct st_image {
	int w;
	int h;
	int c;
	float *data;
} image;

void CuImgResizeMalloc(st_cuda_resize_dev_space & space, int malloc_img_size)
{
	if (space.img_size<malloc_img_size)
	{
		CuImgResizeFree(space);

		space.img_size = malloc_img_size;

		cudaMalloc((void**)&space.img_rgb8, malloc_img_size * sizeof(unsigned char) * 3);
		cudaMalloc((void**)&space.img_rgb16, malloc_img_size * sizeof(unsigned short) * 3);
		cudaMalloc((void**)&space.img_rgbf, malloc_img_size * sizeof(float) * 3);
		cudaMalloc((void**)&space.img_rgbf_part, malloc_img_size * sizeof(float) * 3);
		cudaMalloc((void**)&space.img_rgbf_resized, malloc_img_size * sizeof(float) * 3);
	}
}

void CuImgResizeFree(st_cuda_resize_dev_space &space)
{
	cudaFree(space.img_rgb8);
	cudaFree(space.img_rgb16);
	cudaFree(space.img_rgbf);
	cudaFree(space.img_rgbf_part);
	cudaFree(space.img_rgbf_resized);
}

__device__ float get_pixel(image m, int x, int y, int c)
{
	//assert(x < m.w && y < m.h && c < m.c);
	return m.data[c*m.h*m.w + y*m.w + x];
}

__device__ void set_pixel(image m, int x, int y, int c, float val)
{
	if (x < 0 || y < 0 || c < 0 || x >= m.w || y >= m.h || c >= m.c) return;
	//assert(x < m.w && y < m.h && c < m.c);
	m.data[c*m.h*m.w + y*m.w + x] = val;
}
__device__ void add_pixel(image m, int x, int y, int c, float val)
{
	//assert(x < m.w && y < m.h && c < m.c);
	m.data[c*m.h*m.w + y*m.w + x] += val;
}

__global__ void KernelResizePart1(st_image im, st_image im_part,int w, int h)
{
	unsigned int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	if (threadId >= im_part.w*im_part.h)
	{
		return;
	}

	int r, c, k;
	float w_scale = (float)(im.w - 1) / (w - 1);
	//float h_scale = (float)(im.h - 1) / (h - 1);
	
	r = threadId / im_part.w;
	c = threadId - r*im_part.w;

	for (k = 0; k < im.c; ++k)
	{
		float val = 0;
		if (c == w - 1 || im.w == 1)
		{
			val = get_pixel(im, im.w - 1, r, k);
		}
		else
		{
			float sx = c*w_scale;
			int ix = (int)sx;
			float dx = sx - ix;
			val = (1 - dx) * get_pixel(im, ix, r, k) + dx * get_pixel(im, ix + 1, r, k);
		}
		set_pixel(im_part, c, r, k, val);
	}
}

__global__ void KernelResizePart2(st_image im,st_image im_part, st_image im_resize, int w, int h)
{
	unsigned int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	if (threadId >= im_resize.w*im_resize.h)
	{
		return;
	}

	int r, c, k;
	//float w_scale = (float)(im.w - 1) / (w - 1);
	float h_scale = (float)(im.h - 1) / (h - 1);

	r = threadId / im_resize.w;
	c = threadId - r*im_resize.w;

//	int total_size = w*h;
	for (k = 0; k < im.c; ++k)
	{
		float sy = r*h_scale;
		int iy = (int)sy;
		float dy = sy - iy;
		
		float val = (1 - dy) * get_pixel(im_part, c, iy, k);
		set_pixel(im_resize, c, r, k, val);

		if (r == h - 1 || im.h == 1)
		{
			val = get_pixel(im_part, c, iy, k);
			set_pixel(im_resize, c, h - 1, k, val);	
		}
		else
		{
			val = dy * get_pixel(im_part, c, iy + 1, k);
			add_pixel(im_resize, c, r, k, val);
		}
	}
}

__global__ void KernelInt8ToFloat(unsigned char * img_rgb8,float * img_t, int w, int h)
{
	unsigned int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	unsigned int img_size = w*h;
	if (threadId >= img_size)
	{
		return;
	}
	img_t[threadId]					= float(img_rgb8[threadId]) / 255.f;
	img_t[threadId + img_size]		= float(img_rgb8[threadId + img_size]) / 255.f;
	img_t[threadId + img_size * 2]	= float(img_rgb8[threadId + img_size * 2]) / 255.f;
}

__global__ void KernelInt16ToFloat(unsigned short * img_rgb16, float * img_t, int w, int h)
{
	unsigned int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	unsigned int img_size = w*h;
	if (threadId >= img_size)
	{
		return;
	}
	img_t[threadId]					= float(img_rgb16[threadId]) / 65535.f;
	img_t[threadId + img_size]		= float(img_rgb16[threadId + img_size]) / 65535.f;
	img_t[threadId + img_size * 2]	= float(img_rgb16[threadId + img_size * 2]) / 65535.f;
}

void CuImgResizeOp(int w, int h, st_cuda_resize_dev_space &space, int w_r, int h_r)
{
	st_image im_orig, im_part, im_resize;

	im_orig.w = w;
	im_orig.h = h;
	im_orig.c = 3;
	im_orig.data = space.img_rgbf;

	im_part.w = w_r;
	im_part.h = h;
	im_part.c = 3;
	im_part.data = space.img_rgbf_part;

	{
		dim3 threads = dim3(CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		dim3 blocks = dim3((w_r*h + CUDA_IMG_RESIZE_THREAD_NUM) / CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		KernelResizePart1 << <blocks, threads >> >(im_orig, im_part, w_r, h_r);
	}
	//cudaDeviceSynchronize();

	im_resize.w = w_r;
	im_resize.h = h_r;
	im_resize.c = 3;
	im_resize.data = space.img_rgbf_resized;

	{
		dim3 threads = dim3(CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		dim3 blocks = dim3((w_r*h_r + CUDA_IMG_RESIZE_THREAD_NUM) / CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		KernelResizePart2 << <blocks, threads >> >(im_orig, im_part, im_resize, w_r, h_r);
	}
	//cudaDeviceSynchronize();
}

__global__ void KernelCopyMakeBorder(float * im_resize, int resize_w,int resize_h, float * im_letterbox, int lb_w, int lb_h, int left, int top)
{
	unsigned int threadId = blockIdx.x * CUDA_IMG_RESIZE_THREAD_NUM + threadIdx.x;
	unsigned int img_size = lb_w*lb_h;
	if (threadId >= img_size)
	{
		return;
	}
	int x = threadId%lb_w - left;
	int y = threadId/lb_w - top;
	int resize_img_size = resize_w*resize_h;

	if (x >= 0 && x<resize_w &&y >= 0 && y<resize_h)
	{
		int index = x + y*resize_w;

		im_letterbox[threadId] = im_resize[index];
		threadId += img_size; index += resize_img_size;

		im_letterbox[threadId] = im_resize[index];
		threadId += img_size; index += resize_img_size;

		im_letterbox[threadId] = im_resize[index];
	}
	else
	{
		im_letterbox[threadId] = 114.f / 255.f;
		threadId += img_size; 
		im_letterbox[threadId] = 114.f / 255.f;
		threadId += img_size; 
		im_letterbox[threadId] = 114.f / 255.f;
	}
}

void CuImgCopyMakeBorder(int resize_w, int resize_h,st_cuda_resize_dev_space &space, int net_w, int net_h)
{
	st_image im_resize, im_letterbox;
	im_resize.w = resize_w;
	im_resize.h = resize_h;
	im_resize.c = 3;
	im_resize.data = space.img_rgbf_resized;

	im_letterbox.w = net_w;
	im_letterbox.h = net_h;
	im_letterbox.c = 3;
	im_letterbox.data = space.img_rgbf_resized_letterbox;

	float padd_w = (net_w - resize_w) / 2;
	float padd_h = (net_h - resize_h) / 2;
	float top = padd_h;
	float left = padd_w;
	{
		dim3 threads = dim3(CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		dim3 blocks = dim3((net_w*net_h + CUDA_IMG_RESIZE_THREAD_NUM) / CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		KernelCopyMakeBorder << <blocks, threads >> >(
			im_resize.data, im_resize.w, im_resize.h,
			im_letterbox.data, im_letterbox.w, im_letterbox.h,
			left, top);
	}
}

void CuImgResizeLetterBox(int w, int h, st_cuda_resize_dev_space &space, int net_w, int net_h)
{
	float ratio_w = float(net_w) / float(w), ratio_h = float(net_h) / float(h);
	float ratio = __min(ratio_w, ratio_h);
	int resize_w = ratio_w > ratio_h ? ratio * w : net_w;
	int resize_h = ratio_w > ratio_h ? net_h : ratio * h;
	CuImgResizeOp(w, h, space, resize_w, resize_h);
	CuImgCopyMakeBorder(resize_w, resize_h, space, net_w, net_h);
}


void CuImgResizeRgb8(unsigned char ** host_img_rgb8,unsigned char ** dev_img_rgb8, int w, int h, st_cuda_resize_dev_space &space, int net_w, int net_h)
{
	cudaError_t error; //store cuda error codes
	for (int i = 0; i < 3; ++i)
	{
		if(host_img_rgb8)
		{
			error = cudaMemcpy(space.img_rgb8 + i*w*h, host_img_rgb8[i], w*h * sizeof(unsigned char), cudaMemcpyHostToDevice);
		}
		if(dev_img_rgb8)
		{
			error = cudaMemcpy(space.img_rgb8 + i*w*h, dev_img_rgb8[i], w*h * sizeof(unsigned char), cudaMemcpyDeviceToDevice);
		}
#ifdef _DEBUG
		if (error != cudaSuccess)
		{
			printf("cudaMemcpy, returned error %s (code %d), line(%d)\n", cudaGetErrorString(error), error, __LINE__);
			exit(EXIT_FAILURE);
		}
#endif
	}
	{
		dim3 threads = dim3(CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		dim3 blocks = dim3((w*h + CUDA_IMG_RESIZE_THREAD_NUM) / CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		KernelInt8ToFloat << <blocks, threads >> >(space.img_rgb8, space.img_rgbf, w, h);
		//cudaDeviceSynchronize();
	}
	CuImgResizeLetterBox(w, h, space, net_w, net_h);
}


void CuImgResizeRgb16(unsigned short ** host_img_rgb16,unsigned short ** dev_img_rgb16,int w, int h, st_cuda_resize_dev_space &space, int net_w, int net_h)
{
	cudaError_t error; //store cuda error codes
	for (int i = 0; i < 3; ++i)
	{
		if(host_img_rgb16)
		{
			error = cudaMemcpy(space.img_rgb16 + i*w*h, host_img_rgb16[i], w*h * sizeof(unsigned short), cudaMemcpyHostToDevice);
		}
		if(dev_img_rgb16)
		{
			error = cudaMemcpy(space.img_rgb16 + i*w*h, dev_img_rgb16[i], w*h * sizeof(unsigned short), cudaMemcpyDeviceToDevice);
		}
#ifdef _DEBUG
		if (error != cudaSuccess)
		{
			printf("cudaMemcpy, returned error %s (code %d), line(%d)\n", cudaGetErrorString(error), error, __LINE__);
			exit(EXIT_FAILURE);
		}
#endif
	}
	{
		dim3 threads = dim3(CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		dim3 blocks = dim3((w*h + CUDA_IMG_RESIZE_THREAD_NUM) / CUDA_IMG_RESIZE_THREAD_NUM, 1, 1);
		KernelInt16ToFloat << <blocks, threads >> >(space.img_rgb16, space.img_rgbf, w, h);
	}

	CuImgResizeLetterBox(w, h, space, net_w, net_h);
}

void CuImgResizeRgbF(float** host_img_rgbf,float** dev_img_rgbf, int w, int h, st_cuda_resize_dev_space &space, int net_w, int net_h)
{
	cudaError_t error; //store cuda error codes
 
	for (int i = 0; i < 3; ++i)
	{
		if(host_img_rgbf)
		{
			error = cudaMemcpy(space.img_rgbf + i*w*h, host_img_rgbf[i], w*h * sizeof(float), cudaMemcpyHostToDevice);
		}
		if(dev_img_rgbf)
		{
			error = cudaMemcpy(space.img_rgbf + i*w*h, dev_img_rgbf[i], w*h * sizeof(float), cudaMemcpyDeviceToDevice);
		}
#ifdef _DEBUG
		if (error != cudaSuccess)
		{
			printf("cudaMemcpy, returned error %s (code %d), line(%d)\n", cudaGetErrorString(error), error, __LINE__);
			exit(EXIT_FAILURE);
		}
#endif
	}
	CuImgResizeLetterBox(w, h, space, net_w, net_h);
}